# import paddle.fluid.layers as L
# from paddle.fluid.initializer import Normal
# from paddle.fluid.dygraph import Layer, LayerList, Sequential, Linear, ParameterList
#
# from helm.dynamic.models.nas.nasbench201.operations import OPS, ReLUConvBN
# from helm.dynamic.models.layers import BN, GlobalAvgPool, Conv2d, Pool2d
#
#
# def set_trainable(m, flag):
#     for p in m.parameters():
#         if 'batch_norm' in p.name:
#             pass
#         else:
#             p.trainable = flag
#
#
# class MixedOp(Layer):
#
#     def __init__(self, primitives, C):
#         super().__init__()
#         self._ops = LayerList()
#         for primitive in primitives:
#             op = OPS[primitive](C, 1)
#             if 'pool' in primitive:
#                 op = Sequential(op, BN(C))
#             self._ops.append(op)
#
#     def forward(self, x, hardwts, index):
#         return sum(hardwts[i] * op(x) if i == index else hardwts[i] for i, op in enumerate(self._ops))
#
#
# class ReductionCell(Layer):
#
#     def __init__(self, in_channels, out_channels):
#         super().__init__()
#
#         self.branch1 = Sequential(
#             ReLUConvBN(in_channels, out_channels, 3, stride=2),
#             ReLUConvBN(out_channels, out_channels, 3)
#         )
#         self.branch2 = Sequential(
#             Pool2d(2, stride=2, type='avg'),
#             ReLUConvBN(in_channels, out_channels, 1)
#         )
#
#     def forward(self, x):
#         return self.branch1(x) + self.branch2(x)
#
#
# class NormalCell(Layer):
#
#     def __init__(self, primitives, nodes, C):
#         super().__init__()
#
#         self.nodes = nodes
#
#         self.ops = LayerList()
#         for i in range(1, self.nodes):
#             ops = LayerList()
#             for j in range(nodes):
#                 op = MixedOp(primitives, C)
#                 ops.append(op)
#             self.ops.append(ops)
#
#     def forward(self, x, hardwts, indices):
#         states = [x]
#         for i in range(self.nodes - 1):
#             s = sum(self.ops[i][j](h, hardwts[i][j], int(indices[i][j])) for j, h in enumerate(states))
#             states.append(s)
#         return states[-1]
#
#
# def gumbel_sample(a, tau=1.0):
#     o = -L.log(-L.log(L.uniform_random(a.shape, min=0.0, max=1.0)))
#     return L.softmax((L.log_softmax(a, axis=1) + o) / tau, axis=1)
#
#
# class Network(Layer):
#
#     def __init__(self, primitives=PRIMITIVES, C=16, num_stacked=5, nodes=4, num_classes=10, tau=10.0):
#         super().__init__()
#         self.primitives = primitives
#         self.C = C
#         self.num_classes = num_classes
#         self.num_stacked = num_stacked
#         self.nodes = nodes
#         self.tau = tau
#
#         self.stem = Conv2d(3, C, kernel_size=3, bn=True)
#         for i in range(3):
#             if i != 0:
#                 self.add_sublayer("reduce%d" % i, ReductionCell(C, C * 2))
#                 C = C * 2
#             stage = LayerList()
#             for _ in range(num_stacked):
#                 stage.append(NormalCell(primitives, nodes, C))
#             self.add_sublayer("stage%d" % (i + 1), stage)
#
#         self.avg_pool = GlobalAvgPool()
#         self.classifier = Linear(C, num_classes)
#
#         self._initialize_alphas()
#
#     def forward(self, x):
#         x = self.stem(x)
#         weights = [gumbel_sample(w, self.tau) for w in self.alphas]
#         indices = [L.argmax(w, axis=1) for w in weights]
#         one_h = [L.one_hot(L.unsqueeze(i, 1), len(self.primitives)) for i in indices]
#         indices = [index.numpy() for index in indices]
#         hardwts = [h - w.detach() + w for w, h in zip(weights, one_h)]
#         for cell in self.stage1:
#             x = cell(x, hardwts, indices)
#         x = self.reduce1(x)
#         for cell in self.stage2:
#             x = cell(x, hardwts, indices)
#         x = self.reduce2(x)
#         for cell in self.stage3:
#             x = cell(x, hardwts, indices)
#         x = self.avg_pool(x)
#         x = self.classifier(x)
#         return x
#
#     def _initialize_alphas(self):
#         num_ops = len(self.primitives)
#
#         self.alphas = ParameterList([
#             L.create_parameter([i, num_ops], 'float32', default_initializer=Normal(0, 1e-3))
#             for i in range(1, self.nodes)
#         ])
#
#     def model_parameters(self):
#         ids = set(id(p) for p in self.arch_parameters())
#         for p in self.parameters():
#             if id(p) not in ids:
#                 yield p
#
#     def arch_parameters(self):
#         return self.alphas.parameters()
#
#
# def test_net():
#     net = Network()
